from .module import PREPROCESS
import cv2
import numpy as np


@PREPROCESS.register_module
class Tencrop:

    def __init__(self, size):
        self.size = size

    def process(self, img, anno):
        anno = anno.copy()
        h, w = self.size
        img_h, img_w, c = img.shape
        if img_h < h or img_w < w:
            print('not big enough for tencrop')
            return cv2.resize(img, dsize=tuple(self.size)), anno

        ten_img = np.zeros(shape=tuple(self.size + [10, c]), dtype=img.dtype)
        five_points = [(0, 0), (img_h - h, 0), (0, img_w - w), (img_h - h, img_w - w), ((img_h - h) // 2, (img_w - w) // 2)]
        for idx, (c_y, c_x) in enumerate(five_points):
            temp = img[c_y:c_y + h, c_x:c_x + w][:, :, np.newaxis]
            ten_img[:, :, idx * 2:idx * 2 + 1, :] = temp
            temp = cv2.flip(img[c_y:c_y + h, c_x:c_x + w], 1)[:, :, np.newaxis]
            ten_img[:, :, idx * 2 + 1:idx * 2 + 2, :] = temp
        img = ten_img.reshape(tuple(self.size + [10 * c]))
        return img, anno
